import numpy as np 
from Parameters_1D import *
from scipy import integrate
from model_var_disc_vectorized import *
from model_simulation import simulate_Richards_equation
soilPars_mz1 = pars_mz1()
soilPars_mz2 = pars_mz2()
soilPars_mz3 = pars_mz3()
totalDepth, axialNodes, totalNodes= spatialVariables_1D_act()
#samplingTime, samplingTimeInternal, internalTimeSteps, totTimeSteps = temporalVariable()

def compute_root_zone_moisture(Xk, rooting_depth):
    if rooting_depth == 1.00:    
        root_zone_head = 0.10*((1/6)*(Xk[0] + Xk[1] + Xk[2] + Xk[3] + Xk[4] + Xk[5]))  + \
                         0.20*((1/6)*(Xk[5] + Xk[6] + Xk[7] + Xk[8] + Xk[9] + Xk[10])) + \
                         0.30*((1/11)*(Xk[10] + Xk[11] + Xk[12] + Xk[13] + Xk[14] + Xk[15] + Xk[16] + Xk[17] + Xk[18] + Xk[19] + Xk[20])) + \
                         0.40*((1/11)*(Xk[20] + Xk[21] + Xk[22] + Xk[23] + Xk[24] + Xk[25] + Xk[26] + Xk[27] + Xk[28] + Xk[29] + Xk[30]))
    
    elif rooting_depth == 0.75:
        root_zone_head = 0.10*(1/5)*(Xk[5]  +  Xk[6]  +  Xk[7]  +  Xk[8] + (0.25*(Xk[9] + Xk[8])))  + \
                         0.20*(1/8)*(0.25*(Xk[9] + Xk[8]) + Xk[9] + Xk[10] +  Xk[11] +  Xk[12] +  Xk[13] +  Xk[14] + Xk[15])  + \
                         0.30*(1/9)*(Xk[15] + Xk[16]  + Xk[17]  +  Xk[18] +  Xk[19] + Xk[20] + Xk[21] + Xk[22] + 0.5*(Xk[22] + Xk[23])) + \
                         0.40*(1/9)*(0.5*(Xk[22] + Xk[23]) + Xk[23] + Xk[24] + Xk[25] + Xk[26] + Xk[27] + Xk[28] + Xk[29] + Xk[30])
                         
    elif rooting_depth == 0.50:
        root_zone_head = 0.10*((1/6)*(Xk[10] + Xk[11] + Xk[12] + Xk[13] + Xk[14] + Xk[15])) +\
                         0.20*((1/6)*(Xk[15] + Xk[16] + Xk[17] + Xk[18] + Xk[19] + Xk[20])) +\
                         0.30*((1/6)*(Xk[20] + Xk[21] + Xk[22] + Xk[23] + Xk[24] + Xk[25])) + \
                         0.40*((1/6)*(Xk[25] + Xk[26] + Xk[27] + Xk[28] + Xk[29] + Xk[30]))
    
    else:
        root_zone_head = 0.10*(1/4)*(0.5*(Xk[22]+Xk[23]) + Xk[22] + Xk[21] + Xk[20])   +  \
                         0.20*(1/4)*(0.5*(Xk[22]+Xk[23]) + Xk[23] + Xk[24] + Xk[25])   + \
                         0.30*(1/4)*(Xk[25] + Xk[26] + Xk[27] + 0.5*(Xk[27] + Xk[28])) + \
                         0.40*(1/4)*(0.5*(Xk[27] + Xk[28]) + Xk[28] + Xk[29] + Xk[30])
        
    return root_zone_head                                                                                                                                                                             


def obtain_soil_moisture(head, soilPars):
    vol_moist = volMoistureAllNodes(head, soilPars)
    return vol_moist

#Management Zone 1
def SimulateTheActualPlant_mz1(initialCondition, prescribedIrrigation, referenceEvap, cropCoeff, sim_period, rooting_depth, lai_factor):
    irrigationSequence = prescribedIrrigation
    cropCoeffSequence = cropCoeff
    refEvapSequence = referenceEvap
    headArray= np.zeros((sim_period+1, totalNodes))
    headArray_rz = np.zeros(sim_period+1)
    headArray[0,:] = initialCondition
    headArray_rz[0] = compute_root_zone_moisture(obtain_soil_moisture(initialCondition, soilPars_mz1), rooting_depth[0])
    i=1
    headArray[i,:] = simulate_Richards_equation(headArray[i-1], irrigationSequence[i-1], refEvapSequence[i-1],cropCoeffSequence[i-1],4, soilPars_mz1, rooting_depth[i-1], lai_factor[i-1])
    headArray_rz[i] = compute_root_zone_moisture(obtain_soil_moisture(headArray[i,:], soilPars_mz1), rooting_depth[i-1])

    for i in range(2, sim_period + 1):
        print(i)
        headArray[i,:] = simulate_Richards_equation(headArray[i-1], irrigationSequence[i-1], refEvapSequence[i-1],cropCoeffSequence[i-1],4, soilPars_mz1, rooting_depth[i-1], lai_factor[i-1])
        headArray_rz[i] = compute_root_zone_moisture(obtain_soil_moisture(headArray[i,:], soilPars_mz1), rooting_depth[i-1])
        pass
    return headArray, headArray_rz

#Management Zone 2
def SimulateTheActualPlant_mz2(initialCondition, prescribedIrrigation, referenceEvap, cropCoeff,sim_period, rooting_depth, lai_factor):
    irrigationSequence = prescribedIrrigation
    cropCoeffSequence = cropCoeff
    refEvapSequence = referenceEvap
    headArray= np.zeros((sim_period+1, totalNodes))
    headArray_rz = np.zeros(sim_period+1)
    headArray[0,:] = initialCondition
    headArray_rz[0] = compute_root_zone_moisture(obtain_soil_moisture(initialCondition, soilPars_mz2), rooting_depth[0])
    i=1
    headArray[i,:] = simulate_Richards_equation(headArray[i-1], irrigationSequence[i-1], refEvapSequence[i-1],cropCoeffSequence[i-1],4, soilPars_mz2, rooting_depth[i-1], lai_factor[i-1])
    headArray_rz[i] = compute_root_zone_moisture(obtain_soil_moisture(headArray[i,:], soilPars_mz2), rooting_depth[i-1])

    for i in range(2, sim_period + 1):
        print(i)
        headArray[i,:] = simulate_Richards_equation(headArray[i-1], irrigationSequence[i-1], refEvapSequence[i-1],cropCoeffSequence[i-1],4, soilPars_mz2, rooting_depth[i-1], lai_factor[i-1])
        headArray_rz[i] = compute_root_zone_moisture(obtain_soil_moisture(headArray[i,:], soilPars_mz2), rooting_depth[i-1])
        pass
    return headArray, headArray_rz

#Management Zone 3
def SimulateTheActualPlant_mz3(initialCondition, prescribedIrrigation, referenceEvap, cropCoeff,sim_period, rooting_depth, lai_factor):
    irrigationSequence = prescribedIrrigation
    cropCoeffSequence = cropCoeff
    refEvapSequence = referenceEvap
    headArray= np.zeros((sim_period+1, totalNodes))
    headArray_rz = np.zeros(sim_period+1)
    headArray[0,:] = initialCondition
    headArray_rz[0] = compute_root_zone_moisture(obtain_soil_moisture(initialCondition, soilPars_mz3), rooting_depth[0])
    i=1
    headArray[i,:] = simulate_Richards_equation(headArray[i-1], irrigationSequence[i-1], refEvapSequence[i-1],cropCoeffSequence[i-1],4, soilPars_mz3, rooting_depth[i-1], lai_factor[i-1])
    headArray_rz[i] =  compute_root_zone_moisture(obtain_soil_moisture(headArray[i,:], soilPars_mz3), rooting_depth[i-1])

    for i in range(2, sim_period + 1):
        print(i)
        headArray[i,:] = simulate_Richards_equation(headArray[i-1], irrigationSequence[i-1], refEvapSequence[i-1],cropCoeffSequence[i-1],4, soilPars_mz3, rooting_depth[i-1], lai_factor[i-1])
        headArray_rz[i] = compute_root_zone_moisture(obtain_soil_moisture(headArray[i,:], soilPars_mz3), rooting_depth[i-1])
        pass
    return headArray, headArray_rz
